import argparse

import dataset
import torch
import torch.nn.functional as F
import torch.optim as optim
from torch import nn

from iflearner.business.homo import train_client
from iflearner.business.homo.argument import parser


class Mnist(train_client.PyTorchTrainer):
    class Model(nn.Module):
        def __init__(self, num_channels, num_classes):
            super().__init__()
            self.conv1 = nn.Conv2d(num_channels, 10, kernel_size=5)
            self.conv2 = nn.Conv2d(10, 20, kernel_size=5)
            self.conv2_drop = nn.Dropout2d()
            self.fc1 = nn.Linear(320, 50)
            self.fc2 = nn.Linear(50, num_classes)

        def forward(self, x):
            x = F.relu(F.max_pool2d(self.conv1(x), 2))
            x = F.relu(F.max_pool2d(self.conv2_drop(self.conv2(x)), 2))
            x = x.view(-1, x.shape[1] * x.shape[2] * x.shape[3])
            x = F.relu(self.fc1(x))
            x = F.dropout(x, training=self.training)
            x = self.fc2(x)
            return F.log_softmax(x, dim=1)

    def __init__(
        self, num_clients: int, client_idx: int, iid: bool, lr=0.15, momentum=0.5
    ) -> None:
        self._lr = lr
        self._device = (
            torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu")
        )
        print(f"device: {self._device}")
        self._model = self.Model(num_channels=1, num_classes=10).to(self._device)

        super().__init__(self._model)

        self._optimizer = optim.SGD(self._model.parameters(), lr=lr, momentum=momentum)
        self._loss = F.nll_loss
        self._train_data, self._test_data, self._valid_data = dataset.get_data(
            num_clients=num_clients, client_idx=client_idx, iid=iid
        )

    def config(self) -> dict():
        return {"learning_rate": self._lr, "batch_size": len(self._train_data)}

    def fit(self, epoch):
        self._model.to(self._device)
        self._model.train()
        print(
            f"Epoch: {epoch}, the size of training dataset: {len(self._train_data.dataset)}, batch size: {len(self._train_data)}"
        )
        for batch_idx, (data, target) in enumerate(self._train_data):
            data, target = data.to(self._device), target.to(self._device)
            self._optimizer.zero_grad()
            output = self._model(data)
            loss = self._loss(output, target)
            loss.backward()

            # g = yield self.get(self.ParameterType.ParameterGradient)
            # self.set(homo_pb2.AggregateResult(parameters=g.parameters), self.ParameterType.ParameterGradient)

            self._optimizer.step()

    def evaluate(self, epoch):
        self._model.to(self._device)
        self._model.eval()
        test_loss = 0
        correct = 0
        print(f"The size of testing dataset: {len(self._test_data.dataset)}")
        with torch.no_grad():
            for data, target in self._test_data:
                data, target = data.to(self._device), target.to(self._device)
                output = self._model(data)
                test_loss += self._loss(
                    output, target, reduction="sum"
                ).item()  # sum up batch loss
                pred = output.argmax(
                    dim=1, keepdim=True
                )  # get the index of the max log-probability
                correct += pred.eq(target.view_as(pred)).sum().item()
        test_loss /= len(self._test_data.dataset)

        print(
            "Test set: Average loss: {:.4f}, Accuracy: {}/{} ({:.2f}%)".format(
                test_loss,
                correct,
                len(self._test_data.dataset),
                100.0 * correct / len(self._test_data.dataset),
            )
        )


if __name__ == "__main__":
    parser.add_argument("-a", "--all", help="The number of all clients.")
    parser.add_argument("-i", "--index", help="The index of clients.")
    parser.add_argument("-iid", "--iid", help="Use iid (0 | 1).", default=1, type=int)

    args = parser.parse_args()
    print(args)
    mnist = Mnist(int(args.all), int(args.index), bool(args.iid))
    controller = train_client.Controller(args, mnist)
    controller.run()
